import gym
import random

class HeuristicTag(gym.Wrapper):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.pt_action_space = self.action_space[-1]
        self.n_agents = 2
        self.n_preys = 5

    def reset(self, *args, **kwargs):
        obs = super().reset(*args, **kwargs)
        self.last_prey_obs = obs[-1 * self.n_preys: ] # len = n_preys
        return obs[:-1 * self.n_preys] # len = n_agents
    
    def get_action(self, prey_obs):
        predators_pos_relative = prey_obs[2: 2 + 2 * self.n_agents].reshape(self.n_agents, 2)
        dist = (predators_pos_relative ** 2).sum(axis=-1)
        nearest_predator_pos = predators_pos_relative[dist.argmin()]
        x, y = nearest_predator_pos
        if x >= 0 and abs(x) >= abs(y):
            prey_action = 2
        elif x < 0 and abs(x) >= abs(y):
            prey_action = 1
        elif y >= 0 and abs(y) >= abs(x):
            prey_action = 4
        elif y < 0 and abs(y) >= abs(x):
            prey_action = 3
        else:
            prey_action = 0
        
        return prey_action

    def step(self, action):
        
        prey_actions = []
        for i in range(self.n_preys):
            if random.random() < 0.1:
                prey_actions.append(self.pt_action_space.sample())
            else:
                prey_actions.append(self.get_action(self.last_prey_obs[i]))
        
        action = tuple(action) + tuple(prey_actions)
        obs, rew, done, info = super().step(action)
        
        return obs[:-1 * self.n_preys], rew[:-1 * self.n_preys], done[:-1 * self.n_preys], info